﻿-- drop function public.write_fasta(text);
create or replace function public.write_fasta(fasta_qry text, outfile text, width int) returns text AS $$

""" -------------------------[ write_fasta(text) ]-----------------------------

Returns a fasta file given a table where the first two columns are the fasta 
sequence name and the sequence itself. Additional columns are ignored.

Adjust variable conn for appropriate connection parameters

ARGUMENTS 
fasta_qry: 
    a SELECT query returning the sequence name as first column and the sequence 
    as second column. (Additional columns ignored).

outfile: 
    Name for the output file

width:
    Number of bases in each line of sequence.



EXAMPLE:

SELECT write_fasta('SELECT fasta_name, fasta_seq FROM myTable', 80);

Note: 
If the SELECT query contains single quotes, replace with two single quotes.
E.g. SELECT write_fasta('SELECT fasta_name, ''actg'' FROM myTable');

----------------------------------------------------------------------------"""


" --------------------------[ User's input ]--------------------------------- "

# outfile= 'D:/Tritume/write_fasta.fa' ## /tmp/write_fasta.fa

# Maximum number of characters in each fasta line:
# width= 80

" ---------------------------------------------------------------------------- "

import psycopg2
import os
import re

try:
  conn = psycopg2.connect("dbname='Vitelleschi' user='dberaldi' host='localhost' password='mypassword'");
except:
  return("I am unable to connect to the database")

" -----------------------------[ Define functions ]-------------------------- "

def wrapper(line, width):
    """ Wraps the string 'line' to have each line at most 'width' chars long. 
    Returns a list. Used to wrap sequence to FASTA format. E.g.

    >>> wrapper('ABCDEFGHIJK', 3)
    ['ABC', 'DEF', 'GHI', 'JK']    """
    
    wrapped= []
    count= 0
    linesplit= []
    for n in line:
        if count < width-1:
            linesplit.append(n)
            count += 1
        else:
            linesplit.append(n)
            linesplit= ''.join(linesplit)
            wrapped.append(linesplit)
            count= 0
            linesplit= []
    linesplit= ''.join(linesplit)
    wrapped.append(linesplit)
    if wrapped[-1] == '':
        del(wrapped[-1])
    return(wrapped)

" ---------------------------------------------------------------------------- "

cur = conn.cursor()
cur.execute(fasta_qry)

fout= open(outfile, 'w')

row = cur.fetchone()
"""
if len(row) != 2:
  return('Error: fasta_qry must return exactly two columns.\n' + str(len(row)) + ' column(s) found.') 
  fout.close()
  stop  
"""

aa= re.compile('[^ACDEFGHIKLMNPQRSTVWYBUXZacdefghiklmnpqrstvwybuxz\n\r]')
warn= ''

while row != None:
  fasta_name= '>' + str(row[0])
  fasta_seq= row[1]

  if aa.search(row[1]) != None:
      warn= ('Warning: Non fasta character found in sequence...\n')

  fasta_seq= '\n'.join(wrapper(fasta_seq, width))
      
  fout.write(fasta_name + '\n' + fasta_seq + '\n')
  row= cur.fetchone()

fout.close()

os.system('chmod 777 '+ outfile)  

return(warn + 'Fasta file written to ' + outfile)
    
  $$ LANGUAGE plpythonu;

